//=========================================================================
// Copyright (C) 2025 The C++ Component Model(COMO) Open Source Project
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//=========================================================================

/*
 * gcc -O2 -Wall -o iprtn_multi ip_rtnetlink_multi.c
 *
 * 用法:
 *   sudo ./iprtn_multi eth0 deladd 192.168.1.100/24 192.168.1.200/24
 *   sudo ./iprtn_multi eth0 add    10.0.0.10/24 10.0.0.20/24
 *
 * 参数说明:
 *   <iface>    - 网卡名称 (如 eth0)
 *   deladd/add - 模式: deladd 表示先清空旧 IP, 再配置; add 表示追加配置
 *   <ip/prefix> ... - 一个或多个 IP/前缀, 如 192.168.1.100/24
 *
 */

#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <errno.h>

#include <arpa/inet.h>
#include <net/if.h>
#include <netinet/in.h>
#include <sys/socket.h>
#include <sys/ioctl.h>
#include <linux/rtnetlink.h>
#include <linux/netlink.h>

#include <net/if_arp.h>
#include <netinet/if_ether.h>
#include <netpacket/packet.h>
#include <net/ethernet.h>

#define NL_BUFSIZE 8192

/**
 * 删除接口上的所有 IPv4 地址
 */
static int rtnl_del_all_ipv4(int nl, int ifindex)
{
    char buf[NL_BUFSIZE];
    struct {
        struct nlmsghdr nlh;
        struct ifaddrmsg ifa;
    } req;

    memset(&req, '\0', sizeof(req));
    req.nlh.nlmsg_len   = NLMSG_LENGTH(sizeof(struct ifaddrmsg));
    req.nlh.nlmsg_flags = NLM_F_REQUEST | NLM_F_DUMP;
    req.nlh.nlmsg_type  = RTM_GETADDR;
    req.ifa.ifa_family  = AF_INET;

    if (send(nl, &req, req.nlh.nlmsg_len, 0) < 0) {
        perror("send(RTM_GETADDR)");
        return -1;
    }

    int len;
    while ((len = recv(nl, buf, sizeof(buf), 0)) > 0) {
        struct nlmsghdr *nlh;
        for (nlh = (struct nlmsghdr *)buf;  NLMSG_OK(nlh, len);  nlh = NLMSG_NEXT(nlh, len)) {
            if (nlh->nlmsg_type == NLMSG_DONE) {
                return 0;
            }

            if (nlh->nlmsg_type == NLMSG_ERROR) return -1;

            struct ifaddrmsg *ifa = (struct ifaddrmsg *)NLMSG_DATA(nlh);
            if ((ifa->ifa_family != AF_INET) || (ifa->ifa_index != (__u32)ifindex)) {
                continue;
            }

            struct rtattr *rta = IFA_RTA(ifa);
            int rlen = IFA_PAYLOAD(nlh);
            struct in_addr ip;
            ip.s_addr = 0;

            for (;  RTA_OK(rta, rlen);  rta = RTA_NEXT(rta, rlen)) {
                if (rta->rta_type == IFA_LOCAL) {
                    memcpy(&ip, RTA_DATA(rta), sizeof(ip));
                }
            }
            if (ip.s_addr == 0) {
                continue;
            }

            /* 删除 */
            struct {
                struct nlmsghdr nlh;
                struct ifaddrmsg ifa;
                char buf[256];
            } delreq;

            memset(&delreq, '\0', sizeof(delreq));
            delreq.nlh.nlmsg_len   = NLMSG_LENGTH(sizeof(struct ifaddrmsg));
            delreq.nlh.nlmsg_flags = NLM_F_REQUEST;
            delreq.nlh.nlmsg_type  = RTM_DELADDR;
            delreq.ifa             = *ifa;

            struct rtattr *rta2 = (struct rtattr *)(((char *)&delreq) + NLMSG_ALIGN(delreq.nlh.nlmsg_len));
            rta2->rta_type = IFA_LOCAL;
            rta2->rta_len  = RTA_LENGTH(sizeof(ip));
            memcpy(RTA_DATA(rta2), &ip, sizeof(ip));
            delreq.nlh.nlmsg_len = NLMSG_ALIGN(delreq.nlh.nlmsg_len) + RTA_LENGTH(sizeof(ip));

            if (send(nl, &delreq, delreq.nlh.nlmsg_len, 0) < 0) {
                perror("send(RTM_DELADDR)");
                return -1;
            }
        }
    }
    return 0;
}

/**
 * 添加新地址
 */
int rtnl_add_ip(int nl, int ifindex, const char *ip_str, int prefixlen)
{
    struct {
        struct nlmsghdr nlh;
        struct ifaddrmsg ifa;
        char buf[256];
    } req;

    memset(&req, '\0', sizeof(req));
    req.nlh.nlmsg_len   = NLMSG_LENGTH(sizeof(struct ifaddrmsg));
    req.nlh.nlmsg_flags = NLM_F_REQUEST | NLM_F_CREATE | NLM_F_REPLACE;
    req.nlh.nlmsg_type  = RTM_NEWADDR;

    req.ifa.ifa_family    = AF_INET;
    req.ifa.ifa_prefixlen = prefixlen;
    req.ifa.ifa_scope     = 0;
    req.ifa.ifa_index     = ifindex;

    struct in_addr ip;
    if (inet_pton(AF_INET, ip_str, &ip) <= 0) {
        perror("inet_pton(IP)");
        return -1;
    }

    struct rtattr *rta = (struct rtattr *)(((char *)&req) + NLMSG_ALIGN(req.nlh.nlmsg_len));
    rta->rta_type = IFA_LOCAL;
    rta->rta_len  = RTA_LENGTH(sizeof(ip));
    memcpy(RTA_DATA(rta), &ip, sizeof(ip));
    req.nlh.nlmsg_len = NLMSG_ALIGN(req.nlh.nlmsg_len) + RTA_LENGTH(sizeof(ip));

    rta = (struct rtattr *)(((char *)&req) + NLMSG_ALIGN(req.nlh.nlmsg_len));
    rta->rta_type = IFA_ADDRESS;
    rta->rta_len  = RTA_LENGTH(sizeof(ip));
    memcpy(RTA_DATA(rta), &ip, sizeof(ip));
    req.nlh.nlmsg_len = NLMSG_ALIGN(req.nlh.nlmsg_len) + RTA_LENGTH(sizeof(ip));

    if (send(nl, &req, req.nlh.nlmsg_len, 0) < 0) {
        perror("send(RTM_NEWADDR)");
        return -1;
    }
    return 0;
}

/**
 * 获取 ifindex+MAC
 */
int get_mac(const char *ifname, unsigned char *mac, int *ifindex_out)
{
    int fd = socket(AF_INET, SOCK_DGRAM, 0);
    if (fd < 0) {
        return -1;
    }

    struct ifreq ifr;
    memset(&ifr, '\0', sizeof(ifr));
    strncpy(ifr.ifr_name, ifname, IFNAMSIZ-1);

    *ifindex_out = if_nametoindex(ifname);
    if (0 == *ifindex_out) {
        close(fd);
        return -1;
    }

    if (ioctl(fd, SIOCGIFHWADDR, &ifr) < 0) {
        close(fd);
        return -1;
    }
    memcpy(mac, ifr.ifr_hwaddr.sa_data, ETH_ALEN);

    close(fd);
    return 0;
}

struct arp_pkt {
    struct ether_header eth;
    struct {
        uint16_t htype, ptype;
        uint8_t hlen, plen;
        uint16_t oper;
        uint8_t sha[ETH_ALEN];
        uint8_t spa[4];
        uint8_t tha[ETH_ALEN];
        uint8_t tpa[4];
    } __attribute__((packed)) arp;
} __attribute__((packed));

/**
 * 发送 Gratuitous ARP
 */
int send_gratuitous_arp(int ifindex, const unsigned char *mac, const char *ip_str)
{
    int s = socket(AF_PACKET, SOCK_RAW, htons(ETH_P_ARP));
    if (s < 0) {
        return -1;
    }

    struct in_addr ip;
    if (inet_pton(AF_INET, ip_str, &ip) != 1) {
        close(s);
        return -1;
    }

    struct arp_pkt pkt;
    memset(&pkt, '\0', sizeof(pkt));
    memset(pkt.eth.ether_dhost, 0xFFu, ETH_ALEN);
    memcpy(pkt.eth.ether_shost, mac, ETH_ALEN);
    pkt.eth.ether_type = htons(ETH_P_ARP);

    pkt.arp.htype = htons(ARPHRD_ETHER);
    pkt.arp.ptype = htons(ETH_P_IP);
    pkt.arp.hlen  = ETH_ALEN;
    pkt.arp.plen  = 4;
    pkt.arp.oper  = htons(ARPOP_REQUEST);
    memcpy(pkt.arp.sha, mac, ETH_ALEN);
    memcpy(pkt.arp.spa, &ip, 4);
    memset(pkt.arp.tha, 0, ETH_ALEN);
    memcpy(pkt.arp.tpa, &ip, 4);

    struct sockaddr_ll saddr = {0};
    saddr.sll_family  = AF_PACKET;
    saddr.sll_ifindex = ifindex;
    saddr.sll_halen   = ETH_ALEN;
    memset(saddr.sll_addr, 0xFFu, ETH_ALEN);

    for (int i = 0;  i < 3;  i++) {
        sendto(s, &pkt, sizeof(pkt), 0, (struct sockaddr*)&saddr, sizeof(saddr));
        usleep(200*1000);
    }

    close(s);
    return 0;
}

int iprtn_multi(char *argv)
{
    // 拷贝输入字符串以便安全修改
    char buf[512];
    strncpy(buf, argv, sizeof(buf) - 1);
    buf[sizeof(buf) - 1] = '\0';

    // 分词
    char *tokens[64];
    int count = 0;
    char *tok = strtok(buf, " ");
    while (tok && count < 64) {
        tokens[count++] = tok;
        tok = strtok(NULL, " ");
    }

    if (count < 3) {
        fprintf(stderr, "输入格式错误: 需至少包含 接口、模式 和 一个地址\n");
        return 1;
    }

    const char *ifname = tokens[0];
    const char *mode   = tokens[1];

    unsigned char mac[ETH_ALEN];
    int ifindex;
    if (get_mac(ifname, mac, &ifindex) < 0) {
        return 1;
    }

    int nl = socket(AF_NETLINK, SOCK_RAW, NETLINK_ROUTE);
    if (nl < 0) {
        perror("NETLINK_ROUTE");
        return 1;
    }

    if (strcmp(mode, "del") == 0) {
        rtnl_del_all_ipv4(nl, ifindex);
        close(nl);
        return 0;
    }

    if (strcmp(mode, "deladd") == 0) {
        rtnl_del_all_ipv4(nl, ifindex);
    }

    for (int i = 2;  i < count;  i++) {
        char *arg = tokens[i];
        char *slash = strchr(arg, '/');
        if (! slash) {
            fprintf(stderr, "缺少前缀长度: %s\n", arg);
            continue;
        }

        *slash = '\0';
        const char *ip_str = arg;
        int prefix = atoi(slash + 1);

        if (rtnl_add_ip(nl, ifindex, ip_str, prefix) == 0) {
            printf("添加 %s/%d 成功\n", ip_str, prefix);
            send_gratuitous_arp(ifindex, mac, ip_str);
        }
    }

    close(nl);
    return 0;
}

#if 0
/* ---------------- MAIN ---------------- */
int main(int argc, char *argv[])
{
    if (argc < 4) {
        fprintf(stderr, "用法: %s <iface> <deladd|add> ip1/prefix [ip2/prefix ...]\n", argv[0]);
        return 1;
    }
    const char *ifname = argv[1];
    const char *mode   = argv[2];

    unsigned char mac[ETH_ALEN];
    int ifindex;
    if (get_mac(ifname, mac, &ifindex) < 0) {
        return 1;
    }

    int nl = socket(AF_NETLINK, SOCK_RAW, NETLINK_ROUTE);
    if (nl < 0) {
        perror("NETLINK_ROUTE");
        return 1;
    }

    if (strcmp(mode, "deladd") == 0) {
        rtnl_del_all_ipv4(nl, ifindex);
    }

    for (int i = 3;  i < argc;  i++) {
        char *arg = argv[i];
        char *slash = strchr(arg, '/');
        if (!slash) {
            fprintf(stderr, "缺少前缀长度: %s\n", arg);
            continue;
        }
        *slash = '\0';
        const char *ip_str = arg;
        int prefix = atoi(slash+1);

        if (rtnl_add_ip(nl, ifindex, ip_str, prefix) == 0) {
            printf("添加 %s/%d 成功\n", ip_str, prefix);
            send_gratuitous_arp(ifindex, mac, ip_str);
        }
    }

    close(nl);
    return 0;
}
#endif
